#include <iostream>
#include <fstream>
#include <string>
#include <cstring>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <openssl/md5.h>
using namespace std;

#define MULTICAST_GROUP "224.1.1.1"
#define MULTICAST_PORT 5007
#define BUFFER_SIZE 4096

string calculateFileHash(const string& filePath)
{
    ifstream file(filePath, ios::binary);
    if (!file) throw runtime_error("Failed to open file: " + filePath);

    MD5_CTX md5Context;
    MD5_Init(&md5Context);

    char buffer[BUFFER_SIZE];
    while (file.read(buffer, sizeof(buffer)))
	{
        MD5_Update(&md5Context, buffer, file.gcount());
    }
    MD5_Update(&md5Context, buffer, file.gcount());

    unsigned char hash[MD5_DIGEST_LENGTH];
    MD5_Final(hash, &md5Context);

    string hashStr;
    for (int i = 0; i < MD5_DIGEST_LENGTH; i++)
	{
        char hex[3];
        sprintf(hex, "%02x", hash[i]);
        hashStr += hex;
    }
    return hashStr;
}

void receiveFile()
{
    //创建UDPSocket
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("socket");
        return;
    }

    //允许地址复用
    int reuse = 1;
    if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0)
	{
        perror("setsockopt(SO_REUSEADDR)");
        close(sock);
        return;
    }

    //绑定端口
    sockaddr_in localAddr{};
    localAddr.sin_family = AF_INET;
    localAddr.sin_addr.s_addr = htonl(INADDR_ANY);
    localAddr.sin_port = htons(MULTICAST_PORT);
    if (bind(sock, (sockaddr*)&localAddr, sizeof(localAddr)) < 0)
	{
        perror("bind");
        close(sock);
        return;
    }

    //加入多播组
    ip_mreq multicastRequest{};
    multicastRequest.imr_multiaddr.s_addr = inet_addr(MULTICAST_GROUP);
    multicastRequest.imr_interface.s_addr = htonl(INADDR_ANY);
    if (setsockopt(sock, IPPROTO_IP, IP_ADD_MEMBERSHIP, &multicastRequest, sizeof(multicastRequest)) < 0)
	{
        perror("setsockopt(IP_ADD_MEMBERSHIP)");
        close(sock);
        return;
    }

    cout << "[Receiver] Waiting for file on " << MULTICAST_GROUP << ":" << MULTICAST_PORT << "..." << endl;

    //接收元数据
    char metadataBuffer[BUFFER_SIZE];
    sockaddr_in senderAddr{};
    socklen_t senderAddrLen = sizeof(senderAddr);
    ssize_t metadataLen = recvfrom(sock, metadataBuffer, sizeof(metadataBuffer), 0, (sockaddr*)&senderAddr, &senderAddrLen);
    if (metadataLen < 0)
	{
        perror("recvfrom(metadata)");
        close(sock);
        return;
    }
    metadataBuffer[metadataLen] = '\0';
    string metadata(metadataBuffer);
    
    size_t delim1 = metadata.find('|');
    size_t delim2 = metadata.find('|', delim1 + 1);
    string fileName = metadata.substr(0, delim1);
    size_t fileSize = stoul(metadata.substr(delim1 + 1, delim2 - delim1 - 1));
    string expectedHash = metadata.substr(delim2 + 1);
    
    cout << "[Receiver] Receiving file: " << fileName << " (" << fileSize << " bytes)" << endl;

    // 接收文件内容
    string tempFileName = "temp_" + fileName;
    ofstream outFile(tempFileName, ios::binary);
    if (!outFile)
	{
        cerr << "Failed to create temp file: " << tempFileName << endl;
        close(sock);
        return;
    }

    char dataBuffer[BUFFER_SIZE];
    size_t bytesReceived = 0;
    while (bytesReceived < fileSize)
	{
        ssize_t bytesRead = recvfrom(sock, dataBuffer, sizeof(dataBuffer), 0, (sockaddr*)&senderAddr, &senderAddrLen);
        if (bytesRead < 0)
		{
            perror("recvfrom(data)");
            break;
        }
        outFile.write(dataBuffer, bytesRead);
        bytesReceived += bytesRead;
        cout << "\r[Receiver] Progress: " << bytesReceived << "/" << fileSize << " bytes" << flush;
    }
    cout << "\n[Receiver] File received, verifying hash..." << endl;
    outFile.close();

    //校验哈希
    string actualHash = calculateFileHash(tempFileName);
    if (actualHash == expectedHash)
	{
        rename(tempFileName.c_str(), fileName.c_str());
        cout << "[Receiver] File saved as: " << fileName << endl;
    }
	else
	{
        remove(tempFileName.c_str());
        cerr << "[Receiver] File corrupted: Hash mismatch!" << endl;
    }
    close(sock);
}

int main()
{
    receiveFile();
    return 0;
}